/*  
 * i-OSGi - Tunable Bundle Isolation for OSGi
 * Copyright (C) 2011  Sven Schulz
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package org.iosgi.impl.engine;

import java.math.BigInteger;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

import org.iosgi.util.lang.Longs;
import org.sat4j.core.Vec;
import org.sat4j.core.VecInt;
import org.sat4j.pb.IPBSolver;
import org.sat4j.pb.ObjectiveFunction;
import org.sat4j.pb.PseudoOptDecorator;
import org.sat4j.pb.SolverFactory;
import org.sat4j.specs.ContradictionException;
import org.sat4j.specs.IOptimizationProblem;
import org.sat4j.specs.IVec;
import org.sat4j.specs.IVecInt;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import algorithms.graphColorer.heuristics.GreedyColorer;
import algorithms.graphColorer.heuristics.GreedyStrategy;
import dataStructures.graph.ExtendedUndirectedGraph;
import dataStructures.graph.GraphColoring;

/**
 * @author Sven Schulz
 */
public class PseudoBooleanColorer<V extends Comparable<? super V>, E>
		implements ReportingColorer<V, E> {

	private static final Logger LOGGER = LoggerFactory
			.getLogger(PseudoBooleanColorer.class);

	private static class VertexDegreeComparator<V, E> implements Comparator<V> {

		private final ExtendedUndirectedGraph<V, E> graph;

		public VertexDegreeComparator(final ExtendedUndirectedGraph<V, E> graph) {
			this.graph = graph;
		}

		@Override
		public int compare(V u, V v) {
			return graph.degreeOf(u) - graph.degreeOf(v);
		}

	}

	private final ExtendedUndirectedGraph<V, E> graph;

	public PseudoBooleanColorer(final ExtendedUndirectedGraph<V, E> graph) {
		this.graph = graph;
	}

	@Override
	public GraphColoring<V, E> getColoring() {
		try {
			return this.getColoring(null, Long.MAX_VALUE, TimeUnit.DAYS);
		} catch (TimeoutException te) {
			/* Won't happen. */
			assert false;
			return null;
		}
	}

	@Override
	public GraphColoring<V, E> getColoring(ReportListener<V, E> listener,
			long timeout, TimeUnit unit) throws TimeoutException {
		long deadline = Longs.saturatedAdd(System.currentTimeMillis(),
				TimeUnit.MILLISECONDS.convert(timeout, unit));
		long ref = System.nanoTime();
		GreedyColorer<V, E> precolorer = new GreedyColorer<V, E>(graph,
				new GreedyStrategy<V, E>());
		GraphColoring<V, E> coloring = precolorer.getColoring();
		long elapsed = System.nanoTime() - ref;
		int colors = coloring.getNumberOfColors();
		LOGGER.debug("computed initial coloring with {} colors in {} ns",
				coloring.getNumberOfColors(), elapsed);
		if (listener != null) {
			listener.handle(this, coloring);
		}
		coloring = this.getColoring(colors, listener,
				deadline - System.currentTimeMillis(), TimeUnit.MILLISECONDS);
		return coloring;
	}

	private GraphColoring<V, E> getColoring(int colors,
			ReportListener<V, E> listener, long timeout, TimeUnit unit)
			throws java.util.concurrent.TimeoutException {
		LOGGER.debug("compute coloring with max. {} colors", colors);

		long baseref = System.nanoTime();

		long deadline = Longs.saturatedAdd(System.currentTimeMillis(),
				TimeUnit.MILLISECONDS.convert(timeout, unit));

		/* Create solver. */
		IPBSolver solver = new PseudoOptDecorator(SolverFactory.newDefault());

		try {
			List<V> vertices = new ArrayList<V>(graph.vertexSet());
			Collections.sort(vertices);
			int numberOfVertices = vertices.size();

			/* Indicator variables and variables for tracking used colors. */
			int indicatorVars = numberOfVertices * colors;
			solver.newVar(indicatorVars + colors);

			/* PB constraints to enforce assignment of exactly one color. */
			for (int i = 0; i < numberOfVertices; i++) {
				IVecInt literals = new VecInt(colors);
				IVec<BigInteger> coefficients = new Vec<BigInteger>(colors);
				for (int j = 1; j <= colors; j++) {
					literals.push(i * colors + j);
					coefficients.push(BigInteger.ONE);
				}
				solver.addPseudoBoolean(literals, coefficients, false,
						BigInteger.ONE);
				solver.addPseudoBoolean(literals, coefficients, true,
						BigInteger.ONE);
			}

			/*
			 * CNF constraints enforcing that no adjacent vertices have the same
			 * color.
			 */
			for (int i = 0; i < vertices.size(); i++) {
				V u = vertices.get(i);
				for (int j = 0; j < vertices.size(); j++) {
					V v = vertices.get(j);
					if (graph.containsEdge(u, v)) {
						for (int k = 1; k <= colors; k++) {
							IVecInt literals = new VecInt(2);
							literals.push(-(i * colors + k));
							literals.push(-(j * colors + k));
							solver.addClause(literals);
						}
					}
				}
			}

			/*
			 * CNF constraints (derived from propositional logic expression) to
			 * set track variable, if a color has been assigned to one of the
			 * indicator variables for a given bundle.
			 */
			for (int k = 1; k <= colors; k++) {
				/* First term. */
				IVecInt literals = new VecInt(numberOfVertices + 1);
				literals.push(-(numberOfVertices * colors + k));
				for (int i = 0; i < numberOfVertices; i++) {
					literals.push(i * colors + k);
				}
				solver.addClause(literals);

				/* Second term. */
				for (int i = 0; i < numberOfVertices; i++) {
					literals.clear();
					literals.push(-(i * colors + k));
					literals.push(numberOfVertices * colors + k);
					solver.addClause(literals);
				}
			}

			/* Symmetry breaking. */
			selectiveColoring(colors, solver, vertices);

			IVecInt literals = new VecInt(2);
			for (int k = 1; k <= colors; k++) {
				literals.push(-(numberOfVertices * colors + k));
				for (int i = 1; i < k; i++) {
					literals.push(numberOfVertices * colors + i);
					solver.addClause(literals);
				}
				literals.clear();
			}

			/* Objective function. */
			literals = new VecInt(colors);
			IVec<BigInteger> coefficients = new Vec<BigInteger>(colors);
			for (int k = 1; k <= colors; k++) {
				literals.push(numberOfVertices * colors + k);
				coefficients.push(BigInteger.ONE);
			}
			ObjectiveFunction objFunc = new ObjectiveFunction(literals,
					coefficients);
			solver.setObjectiveFunction(objFunc);

			/* Solve */
			IOptimizationProblem op = (IOptimizationProblem) solver;
			int[] model = null;
			try {
				if (deadline != Long.MAX_VALUE) {
					solver.setTimeoutMs(deadline - System.currentTimeMillis());
				}
				long stepref = System.nanoTime();
				while (System.currentTimeMillis() < deadline
						&& !Thread.currentThread().isInterrupted()
						&& op.admitABetterSolution()) {
					model = solver.model();
					GraphColoring<V, E> coloring = this.getColoring(colors,
							vertices, numberOfVertices, model);
					long now = System.nanoTime();
					LOGGER.debug(
							"solution found with {} colors and objective value {} in {}/{} ns",
							new Object[] { coloring.getNumberOfColors(),
									op.getObjectiveValue(), now - stepref,
									now - baseref });
					if (listener != null) {
						listener.handle(this, coloring);
					}
					op.discardCurrentSolution();
					stepref = System.nanoTime();
					long nextTO = Math.max(0,
							deadline - System.currentTimeMillis());
					solver.setTimeoutMs(nextTO);
					LOGGER.debug("timeout set to {} ms", nextTO);
				}
			} catch (ContradictionException ce) {
				/*
				 * Raised by discardCurrentSolution, but we may have found a
				 * model.
				 */
			} catch (org.sat4j.specs.TimeoutException te) {
				if (model == null) {
					throw new TimeoutException(te.getMessage());
				}
			}
			if (model != null) {
				GraphColoring<V, E> coloring = getColoring(colors, vertices,
						numberOfVertices, model);
				return coloring;
			}
		} catch (ContradictionException ce) {
			LOGGER.error("contradiction encountered when adding constraints",
					ce);
			assert false;
		}
		return null;
	}

	/*
	 * Performs selective coloring as described in
	 * "Breaking Instance-Independent Symmetries in Exact Graph Coloring" by
	 * Ramani et. al.
	 */
	private void selectiveColoring(int colors, IPBSolver solver,
			List<V> vertices) throws ContradictionException {
		VertexDegreeComparator<V, E> c = new VertexDegreeComparator<V, E>(graph);

		/* Precolor vertex v with maximum degree with first color. */
		V v = Collections.max(vertices, c);
		IVecInt literals = new VecInt(1);
		literals.push(vertices.indexOf(v) * colors + 1);
		solver.addClause(literals);
		literals.clear();

		/* Precolor neighbor of v with maximum degree with second color. */
		Collection<V> neighbors = graph.getNeighbors(v);
		if (!neighbors.isEmpty()) {
			v = Collections.max(neighbors, c);
			literals.push(vertices.indexOf(v) * colors + 2);
			solver.addClause(literals);
		}
	}

	private GraphColoring<V, E> getColoring(int colors, List<V> vertices,
			int numberOfVertices, int[] model) {
		GraphColoring<V, E> coloring = new GraphColoring<V, E>(graph);
		Map<Integer, Integer> translationTable = this.createTranslationTable(
				colors, numberOfVertices, model);
		for (int i = 0; i < vertices.size(); i++) {
			V v = vertices.get(i);
			for (int k = 0; k < colors; k++) {
				if (model[i * colors + k] > 0) {
					coloring.setColor(v, translationTable.get(k));
				}
			}
		}
		return coloring;
	}

	/*
	 * As potentially not all colors are used, we have to compactify the color
	 * range.
	 */
	private Map<Integer, Integer> createTranslationTable(int colors,
			int numberOfVertices, int[] model) {
		Map<Integer, Integer> translated = new HashMap<Integer, Integer>();
		int c = 0;
		for (int k = 0; k < colors; k++) {
			if (model[numberOfVertices * colors + k] > 0) {
				translated.put(k, c++);
			}
		}
		return translated;
	}

}